import os
from datetime import datetime
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from langgraph.constants import END, START
from typing import TypedDict, List, Literal
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from src.agentic.agent.multi.StockAgent import StockAgent
from src.agentic.agent.multi.SearchAgent import SearchAgent
from src.agentic.agent.multi.WeatherAgent import WeatherAgent
from src.agentic.agent.multi.CalculateAgent import CalculateAgent

logger = getLogger()

class RouteAct(BaseModel):
    agent: Literal["calculate", "weather", "stock", "search"] = Field(description = "调用执行的下一步节点")

class ReflectState(TypedDict):
    query: str
    documents: List[str]
    answer: str

class ReflectionAgent:

    def __init__(self, llm_model, agent_tools, max_loop: int = 2):
        self.llm_model = llm_model
        self.max_loop = max_loop
        self.agent_tools = {tool.name: tool for tool in agent_tools}

    def generate_node(self, state):
        logger.info("ReflectionAgent generate_node start")
        query = state["query"]
        documents = state.get("documents", None)
        template = """
            当前时间：{current_date}
        
            你是一位精准的问题-答案分析专家，请严格按照要求分析并解答，同时应参考评论建议，判断当前应归属的节点类型。
            
            用户问题： {question}
            评论建议： {documents}
            
            根据用户问题，返回下列节点中的一个：
            calculate: 用于执行数学计算，只能用于数值计算
            weather: 用于查询城市当天及未来的实时天气情况，禁止用来查询过去时间的天气情况
            stock: 用于查询沪深300指数指定股票的开盘/收盘价
            search: 用于搜索真实世界的信息，如事实、人物、地点等，且不属于上述三类时使用
        """
        prompt = ChatPromptTemplate.from_template(template)
        route_chain = prompt | self.llm_model.with_structured_output(RouteAct)
        current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        response = route_chain.invoke({ "current_date": current_date, "question": query, "documents": documents })
        logger.info(f"ReflectionAgent generate_node response: {response}")

        agent = None
        if "calculate" == response.agent:
            agent = CalculateAgent(self.llm_model, self.agent_tools["calculate_numerical"])
        elif "weather" == response.agent:
            agent = WeatherAgent(self.llm_model, self.agent_tools["search_weather"])
        elif "stock" == response.agent:
            agent = StockAgent(self.llm_model, self.agent_tools["snatch_stock_price"])
        elif "search" == response.agent:
            agent = SearchAgent(self.llm_model, self.agent_tools["search_web"])
        response = agent.invoke(query)
        return { "answer": response["answer"] }

    def reflect_node(self, state):
        logger.info("ReflectionAgent reflect_node start")
        query = state["query"]
        answer = state.get("answer", None)
        template = """
            当前时间是：{current_date}，请基于当前时间分析问题。
        
            你是一个问题分析专家，为提供的答案生成评论和建议。请提供详细的建议，包括深度、风格等方面的请求。
            必须使用中文回答。
            
            问题：{question}
            答案：{answer}
            
            评论及建议执行的节点：
            calculate: 用于执行数学计算，只能用于数值计算
            weather: 用于查询城市当天及未来的实时天气情况，禁止用来查询过去时间的天气情况
            stock: 用于查询沪深300指数指定股票的开盘/收盘价
            search: 用于搜索真实世界的信息，如事实、人物、地点等，且不属于上述三类时使用
            
            请用中文尽量详细的分析问题。
        """
        prompt = ChatPromptTemplate.from_template(template)
        reflect_chain = prompt | self.llm_model | StrOutputParser()
        current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        response = reflect_chain.invoke({ "current_date": current_date, "answer": answer, "question": query })
        logger.info(f"ReflectionAgent reflect_node response: \n{response}")
        documents = state.get("documents", [])
        documents.append(response)
        return { "documents": documents }

    def route(self, state):
        logger.info("ReflectionAgent route start")
        documents = state.get("documents", [])
        if len(documents) >= self.max_loop:
            return "__end__"
        else:
            return "reflect"

    def build_graph(self):
        graph = StateGraph(ReflectState)
        graph.add_node("generate", self.generate_node)
        graph.add_node("reflect", self.reflect_node)

        graph.add_edge(START, "generate")
        graph.add_conditional_edges("generate", self.route, { "__end__": END, "reflect": "reflect" })
        graph.add_edge("reflect", "generate")

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agent_Reflection_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"ReflectionAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({"query": query})
        answer = response.get("answer", None)
        logger.info(f"ReflectionAgent invoke answer len: {len(str(answer))}")
        return { "chain_result": answer }
